use sq_global::{API_STOCK_COMP, API_STOCK_DAILY};

use polars::prelude::{
    Column, DataType, GetOutput, IntoColumn, IntoLazy, ListNameSpaceExtension, NamedFrom,
    PolarsResult, RoundMode, SortMultipleOptions, col, lit,
};
use polars::series::Series;
use sq_macro::test_with_ctx;
use sq_polars::{get_col_idx, get_str_col, get_u32_col};
use tokio::try_join;

#[test_with_ctx]
pub async fn test() {
    run().await.expect("run failed");
}

pub async fn run() -> anyhow::Result<Vec<String>> {
    let (com_df, daily_dfs) = try_join!(
        sq_tushare::load_data(API_STOCK_COMP, Vec::new()),
        sq_tushare::recent_dates(API_STOCK_DAILY, 40)
    )?;

    let daily_df = sq_polars::concat_dfs(daily_dfs)?;

    let top10_longest_rise_stock = daily_df
        .sort(["trade_date"], SortMultipleOptions::default())
        .group_by(["ts_code"])
        .agg([
            col("trade_date"),
            col("pct_chg")
                .round(2, RoundMode::HalfAwayFromZero)
                .apply(cont_rise_pct, GetOutput::from_type(DataType::Float64))
                .alias("continuous_chg"),
        ])
        .with_column(
            col("continuous_chg")
                .list()
                .len()
                .alias("continuous_chg_len"),
        )
        .with_columns([
            col("continuous_chg")
                .list()
                .mean()
                .round(2, RoundMode::HalfAwayFromZero)
                .alias("continuous_chg_avg"),
            col("trade_date")
                .list()
                .tail(col("continuous_chg_len"))
                .list()
                .join(lit(","), true)
                .alias("trade_dates"),
            col("continuous_chg")
                .list()
                .eval(col("").cast(DataType::String) + lit("%"), true)
                .list()
                .join(lit(","), true)
                .alias("continuous_chg_fmt"),
        ])
        .inner_join(com_df.lazy(), "ts_code", "ts_code")
        .sort(
            ["continuous_chg_len", "continuous_chg_avg"],
            SortMultipleOptions::default().with_order_descending(true),
        )
        .limit(10)
        .collect()?;

    let map = get_col_idx!(&top10_longest_rise_stock);
    let cm_idx = map["com_name"];
    let tc_idx = map["ts_code"];
    let ccl_idx = map["continuous_chg_len"];
    let td_idx = map["trade_dates"];
    let cc_idx = map["continuous_chg_fmt"];
    let avg_idx = map["continuous_chg_avg"];

    let mut top10 = Vec::new();
    for row_idx in 0..top10_longest_rise_stock.height() {
        let row = top10_longest_rise_stock.get_row(row_idx)?;
        let com_name = get_str_col!(row, cm_idx);
        let com_code = get_str_col!(row, tc_idx);
        let ct_td = get_str_col!(row, td_idx);
        let cont_chg = get_str_col!(row, cc_idx);
        let avg_chg = get_str_col!(row, avg_idx);
        let chg_len = get_u32_col!(row, ccl_idx);

        let msg = format!(
            "上市公司：{com_name}（股票代码{com_code}）最近连续{chg_len}个交易日（{ct_td}）股价上涨，平均涨幅为{avg_chg}%，涨幅分别为{cont_chg}。\n\n"
        );
        top10.push(msg);
    }

    Ok(top10)
}

fn cont_rise_pct(col: Column) -> PolarsResult<Option<Column>> {
    let mut vals: Vec<f64> = col
        .f64()?
        .iter()
        .rev()
        .map_while(|val| val.filter(|&v| v > 0.0))
        .collect();
    vals.reverse();
    Ok(Some(
        Series::new("continuous_chg".into(), vals).into_column(),
    ))
}
